import taichi as ti
from celestial_objects import Star, Planet

ti.init(arch=ti.gpu)

if __name__.__eq__("__main__"):
    window_size = 800
    window = ti.ui.Window('2D星系模拟', (window_size, window_size))
    canvas = window.get_canvas()

    sun = Star(N=1, mass=10000)
    planets = Star(N=3, mass=1000)
    satellites = Planet(N=10000, mass=1)
    all = [sun, planets, satellites]


    def initialize():
        sun.initialize(0.5, 0.5, 0.1, 0)
        planets.initialize(0.5, 0.5, 0.1, 0)
        satellites.initialize(0.5, 0.5, 0.3, 40)


    initialize()
    h = 5e-5
    pause = False
    fileIndex = 0

    while window.running:
        for e in window.get_events(ti.ui.PRESS):
            if e.key == ti.ui.ESCAPE:
                exit(0)
            elif e.key == 'p':
                pause = not pause
                fileIndex = 0
            elif e.key == 'r':
                initialize()
        if pause:
            continue

        sun.clear_force()
        sun.compute_force_self()
        for obj in all:
            if obj != sun:
                sun.compute_force_other(obj)

        planets.clear_force()
        planets.compute_force_self()
        for obj in all:
            if obj != planets:
                planets.compute_force_other(obj)

        satellites.clear_force()
        satellites.compute_force_self()
        for obj in all:
            if obj != satellites:
                satellites.compute_force_other(obj)

        sun.update_pos(h)
        planets.update_pos(h)
        satellites.update_pos(h)

        canvas.circles(satellites.Pos(), radius=1 / window_size, color=(0.5, 0.5, 0.5))
        canvas.circles(planets.Pos(), radius=5 / window_size, color=(0.9, 0.5, 0.8))
        canvas.circles(sun.Pos(), radius=10 / window_size, color=(0.99, 0.99, 0.9))

        window.show()
